import numpy as np
import cv2  # as cv
import tkinter as tk
from tkinter import filedialog
#import os

np.set_printoptions(suppress=True)

root = tk.Tk()
root.geometry("500x500")
root.title("两图片匹配")

img1 = None
img2 = None
fxy1 = 1
fxy2 = 1


def compute_padding(fxy):
    paddingTop = 650
    paddingButtom = 100
    padding_v = padding_e.get()
    padding_arr = padding_v.split(",")
    if padding_arr[0] == '':
        padding_arr = []
    if len(padding_arr) == 1:
        paddingTop = int(padding_arr[0])
    elif len(padding_arr) == 2:
        paddingTop = int(padding_arr[0])
        paddingButtom = int(padding_arr[1])
    #global paddingTop,paddingButtom
    paddingTop *= fxy
    paddingButtom *= fxy
    return paddingTop, paddingButtom


def do_mask_padding(image, paddingTop, paddingButtom):
    maskAll = np.zeros(image.shape[0:2], dtype="uint8")
    cv2.rectangle(maskAll, (0, int(paddingTop)), (image.shape[1], image.shape[0] - int(paddingButtom)), 255, -1)
    gray = cv2.cvtColor(image, cv2.COLOR_RGB2GRAY)
    mask_out = cv2.bitwise_and(gray, maskAll)
    cv2.imwrite('out/orb_mask_out.jpg', mask_out)
    return maskAll


def stitchImage(imageA, imageB):
    # 2.转换为灰度图，
    gray1 = cv2.cvtColor(imageA, cv2.COLOR_BGR2GRAY)
    gray2 = cv2.cvtColor(imageB, cv2.COLOR_BGR2GRAY)

    finder = cv2.ORB_create(nfeatures=500 * 4)#, scaleFactor=2.0)

    paddingTop1,paddingButtom1 = compute_padding(fxy1)
    mask1 = do_mask_padding(img1,paddingTop1,paddingButtom1)
    paddingTop2,paddingButtom2 = compute_padding(fxy2)
    mask2 = do_mask_padding(img2,paddingTop2,paddingButtom2)


    kpsA, despA = finder.detectAndCompute(gray1, mask1)  # 计算关键点和描述符
    kpsB, despB = finder.detectAndCompute(gray2, mask2)

    bf = cv2.BFMatcher(cv2.NORM_HAMMING)  # crossCheck=True knnMatch(k=1)
    matches = bf.knnMatch(despA, despB, 2)
    print(len(matches))
    good_matches = []
    # good_matchesK = []
    for m, n in matches:
        if m.distance < 0.8 * n.distance:  # d1 < 0.4*d2
            good_matches.append(m)
            # good_matchesK.append([m])
            # good_matches.append(m) for drawMatches()
    print(len(good_matches))
    # # Draw top matches,一个点会绘制两条线
    # flags=0表示画特征点和连线，flags=2表示不画特征点
    # imMatches = cv2.drawMatchesKnn(gray1, kpsA, gray2, kpsB, good_matchesK[0:100], None)
    # kpsAgood = []
    # kpsBgood = []
    # for match in good_matches:
    #     kpsAgood.append(kpsA[match.queryIdx])
    #     kpsBgood.append(kpsB[match.trainIdx])
    # imgKp1 = cv2.drawKeypoints(gray1, kpsAgood, None, flags=cv2.DRAW_MATCHES_FLAGS_DRAW_RICH_KEYPOINTS)
    # imgKp2 = cv2.drawKeypoints(gray2, kpsBgood, None, flags=cv2.DRAW_MATCHES_FLAGS_DRAW_RICH_KEYPOINTS)
    imMatches = cv2.drawMatches(gray1, kpsA, gray2, kpsB, good_matches[0:100], None)
    cv2.namedWindow("matches", cv2.WINDOW_FREERATIO)
    cv2.imshow('matches', imMatches)
    cv2.waitKey(0)
    cv2.imwrite("out/matches.jpg", imMatches)
    # result = os.system("explorer out/matches.jpg")
    #result = os.popen("explorer out/matches.jpg\r\n")
    #print(result)
    # import subprocess
    # result = subprocess.run(["explorer", "out/matches.jpg"], stdout=subprocess.PIPE)
    # print(result.stdout.decode())

    # kps1 = np.float32([kp.pt for kp in kpsA])   #求出所有关键点的坐标
    # kps2 = np.float32([kp.pt for kp in kpsB])
    #
    gkps1 = np.float32([kpsA[a.queryIdx].pt for a in good_matches]).reshape(-1, 1, 2)  # 求出可靠匹配的x,y坐标
    gkps2 = np.float32([kpsB[a.trainIdx].pt for a in good_matches]).reshape(-1, 1, 2)
    (M, mask) = cv2.findHomography(gkps2, gkps1, cv2.RANSAC, 4.0)  # 默认3.0，错误阈值,2转换到1的矩阵
    print(M)

    result = cv2.warpPerspective(imageB, M, (imageA.shape[1] + imageB.shape[1], imageB.shape[0]))
    result[0:imageA.shape[0], 0:imageA.shape[1]] = imageA
    return result


def choose_image():
    src_FILEs = filedialog.askopenfilenames(initialdir="images/", title="选择图片",
                                            filetypes=(("JPEG files", "*.jpg"), ("PNG files", "*.png")))
    print(src_FILEs)
    # 1.读入图片
    global img1, img2
    img1 = cv2.imread(src_FILEs[0])  # r"D:\data\IMG_20230719_105059.jpg")
    img2 = cv2.imread(src_FILEs[1])  # r"D:\data\IMG_20230719_105057.jpg")
    #自动缩放图像
    global fxy1, fxy2
    if img1.shape[1] > 1500:
        fxy1 = 1500 / img1.shape[1]
        img1 = cv2.resize(img1, None, fx=fxy1, fy=fxy1)
    if img2.shape[1] > 1500:
        fxy2 = 1500 / img2.shape[1]
        img2 = cv2.resize(img2, None, fx=fxy2, fy=fxy2)


def do_image():
    # 进行拼接
    result = stitchImage(img1, img2)

    cv2.imwrite('out/ORB_knn.jpg', result)
    cv2.namedWindow("result", cv2.WINDOW_FREERATIO)
    cv2.imshow('result', result)

    cv2.waitKey(0)
    cv2.destroyAllWindows()


if __name__ == "__main__":
    button = tk.Button(root, text="选择两张图片", command=choose_image)
    button.pack(side="top", pady=50)

    padding_label = tk.Label(root, text='padding')  # ,anchor='c').grid(row=0,column=0)
    padding_label.pack(side="top")
    padding_e = tk.StringVar()
    padding = tk.Entry(root, width=20, textvariable=padding_e)  # .grid(row=0,column=1)
    padding.pack(side="top")

    button_do = tk.Button(root, text="测试拼接", command=do_image)
    button_do.pack(side="bottom", pady=50)

    out_e = tk.StringVar()
    out_label = tk.Label(root, text='输处目录:', textvariable=out_e)
    out_label.pack(side="bottom")

    root.mainloop()
